{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b906abab",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright 2021 NVIDIA Corporation. All Rights Reserved.\n",
    "#\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "#     http://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License.\n",
    "# =============================================================================="
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb6958f0",
   "metadata": {},
   "source": [
    "# HugeCTR Wide and Deep Model with Criteo\n",
    "\n",
    "## Overview\n",
    "\n",
    "In this notebook, we provide a tutorial that shows how to train a wide and deep model using the high-level Python API from HugeCTR on the original Criteo dataset as training data.\n",
    "We show how to produce prediction results based on different types of local database."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1bf41447",
   "metadata": {},
   "source": [
    "## Dataset Preprocessing\n",
    "\n",
    "### Generate training and validation data folders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ae158c8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# define some data folder to store the original and preprocessed data\n",
    "# Standard Libraries\n",
    "import os\n",
    "from time import time\n",
    "import re\n",
    "import shutil\n",
    "import glob\n",
    "import warnings\n",
    "BASE_DIR = \"/wdl_train\"\n",
    "train_path  = os.path.join(BASE_DIR, \"train\")\n",
    "val_path = os.path.join(BASE_DIR, \"val\")\n",
    "CUDA_VISIBLE_DEVICES = os.environ.get(\"CUDA_VISIBLE_DEVICES\", \"0\")\n",
    "n_workers = len(CUDA_VISIBLE_DEVICES.split(\",\"))\n",
    "frac_size = 0.15\n",
    "allow_multi_gpu = False\n",
    "use_rmm_pool = False\n",
    "max_day = None  # (Optional) -- Limit the dataset to day 0-max_day for debugging\n",
    "\n",
    "if os.path.isdir(train_path):\n",
    "    shutil.rmtree(train_path)\n",
    "os.makedirs(train_path)\n",
    "\n",
    "if os.path.isdir(val_path):\n",
    "    shutil.rmtree(val_path)\n",
    "os.makedirs(val_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9d463e40",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total 14537948\n",
      "-rw-r--r-- 1 root root  3336516608 Jul  5 05:44 0.8870d61b8a1f4deca0f911acfb072999.parquet\n",
      "-rw-r--r-- 1 root root          61 Jul  5 05:44 _file_list.txt\n",
      "-rw-r--r-- 1 root root      602767 Jul  5 05:44 _metadata\n",
      "-rw-r--r-- 1 root root        1538 Jul  5 05:44 _metadata.json\n",
      "drwxr-xr-x 2 root root        4096 Jul  5 05:41 \u001b[0m\u001b[01;34mtemp-parquet-after-conversion\u001b[0m/\n",
      "-rwxrwxr-x 1 1025 1025 11549710546 Jul  5 05:39 \u001b[01;32mtrain.txt\u001b[0m*\n"
     ]
    }
   ],
   "source": [
    "!ls -l $train_path"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99e7db3c",
   "metadata": {},
   "source": [
    "### Download the original Criteo dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6dfa9fa5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reading package lists... Done\n",
      "Building dependency tree       \n",
      "Reading state information... Done\n",
      "wget is already the newest version (1.20.3-1ubuntu1).\n",
      "0 upgraded, 0 newly installed, 0 to remove and 4 not upgraded.\n"
     ]
    }
   ],
   "source": [
    "!apt-get install wget"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0da41ce3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2021-07-05 03:27:50--  http://azuremlsampleexperiments.blob.core.windows.net/criteo/day_0.gz\n",
      "Resolving azuremlsampleexperiments.blob.core.windows.net (azuremlsampleexperiments.blob.core.windows.net)... 20.60.140.36\n",
      "Connecting to azuremlsampleexperiments.blob.core.windows.net (azuremlsampleexperiments.blob.core.windows.net)|20.60.140.36|:80... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 16309554343 (15G) [application/octet-stream]\n",
      "Saving to: ‘wdl_train/train/day_0.gz’\n",
      "\n",
      "day_0.gz               100%[===================================================>]   15G  --.-KB/s    in  6m 12s \n",
      "2021-07-05 03:34:04 (79.2 MB/s) - 'day_0.gz' saved [16309554343/16309554343]"
     ]
    }
   ],
   "source": [
    "!wget -P $train_path http://azuremlsampleexperiments.blob.core.windows.net/criteo/day_0.gz"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4508791b",
   "metadata": {},
   "source": [
    "\n",
    "Split the dataset into training and validation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1f9620ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "!gzip -d -c $train_path/day_0.gz > day_0\n",
    "!head -n 45840617 day_0 > $train_path/train.txt\n",
    "!tail -n 2000000 day_0 > $val_path/test.txt "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9919817c",
   "metadata": {},
   "source": [
    "### Preprocessing with NVTabular"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "616c6f19",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing /wdl_train/preprocess.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile /wdl_train/preprocess.py\n",
    "import os\n",
    "import sys\n",
    "import argparse\n",
    "import glob\n",
    "import time\n",
    "from cudf.io.parquet import ParquetWriter\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import concurrent.futures as cf\n",
    "from concurrent.futures import as_completed\n",
    "import shutil\n",
    "\n",
    "import dask_cudf\n",
    "from dask_cuda import LocalCUDACluster\n",
    "from dask.distributed import Client\n",
    "from dask.utils import parse_bytes\n",
    "from dask.delayed import delayed\n",
    "\n",
    "import cudf\n",
    "import rmm\n",
    "import nvtabular as nvt\n",
    "from nvtabular.io import Shuffle\n",
    "from nvtabular.utils import device_mem_size\n",
    "from nvtabular.ops import Categorify, Clip, FillMissing, HashBucket, LambdaOp, Normalize, Rename, Operator, get_embedding_sizes\n",
    "#%load_ext memory_profiler\n",
    "\n",
    "import logging\n",
    "logging.basicConfig(format='%(asctime)s %(message)s')\n",
    "logging.root.setLevel(logging.NOTSET)\n",
    "logging.getLogger('numba').setLevel(logging.WARNING)\n",
    "logging.getLogger('asyncio').setLevel(logging.WARNING)\n",
    "\n",
    "# define dataset schema\n",
    "CATEGORICAL_COLUMNS=[\"C\" + str(x) for x in range(1, 27)]\n",
    "CONTINUOUS_COLUMNS=[\"I\" + str(x) for x in range(1, 14)]\n",
    "LABEL_COLUMNS = ['label']\n",
    "COLUMNS =  LABEL_COLUMNS + CONTINUOUS_COLUMNS +  CATEGORICAL_COLUMNS\n",
    "#/samples/criteo mode doesn't have dense features\n",
    "criteo_COLUMN=LABEL_COLUMNS +  CATEGORICAL_COLUMNS\n",
    "#For new feature cross columns\n",
    "CROSS_COLUMNS = []\n",
    "\n",
    "\n",
    "NUM_INTEGER_COLUMNS = 13\n",
    "NUM_CATEGORICAL_COLUMNS = 26\n",
    "NUM_TOTAL_COLUMNS = 1 + NUM_INTEGER_COLUMNS + NUM_CATEGORICAL_COLUMNS\n",
    "\n",
    "\n",
    "# Initialize RMM pool on ALL workers\n",
    "def setup_rmm_pool(client, pool_size):\n",
    "    client.run(rmm.reinitialize, pool_allocator=True, initial_pool_size=pool_size)\n",
    "    return None\n",
    "\n",
    "#compute the partition size with GB\n",
    "def bytesto(bytes, to, bsize=1024):\n",
    "    a = {'k' : 1, 'm': 2, 'g' : 3, 't' : 4, 'p' : 5, 'e' : 6 }\n",
    "    r = float(bytes)\n",
    "    return bytes / (bsize ** a[to])\n",
    "\n",
    "class FeatureCross(Operator):\n",
    "    def __init__(self, dependency):\n",
    "        self.dependency = dependency\n",
    "\n",
    "    def transform(self, columns, gdf):\n",
    "        new_df = type(gdf)()\n",
    "        for col in columns.names:\n",
    "            new_df[col] = gdf[col] + gdf[self.dependency]\n",
    "        return new_df\n",
    "\n",
    "    def dependencies(self):\n",
    "        return [self.dependency]\n",
    "\n",
    "#process the data with NVTabular\n",
    "def process_NVT(args):\n",
    "\n",
    "    if args.feature_cross_list:\n",
    "        feature_pairs = [pair.split(\"_\") for pair in args.feature_cross_list.split(\",\")]\n",
    "        for pair in feature_pairs:\n",
    "            CROSS_COLUMNS.append(pair[0]+'_'+pair[1])\n",
    "\n",
    "\n",
    "    logging.info('NVTabular processing')\n",
    "    train_input = os.path.join(args.data_path, \"train/train.txt\")\n",
    "    val_input = os.path.join(args.data_path, \"val/test.txt\")\n",
    "    PREPROCESS_DIR_temp_train = os.path.join(args.out_path, 'train/temp-parquet-after-conversion')\n",
    "    PREPROCESS_DIR_temp_val = os.path.join(args.out_path, 'val/temp-parquet-after-conversion')\n",
    "    PREPROCESS_DIR_temp = [PREPROCESS_DIR_temp_train, PREPROCESS_DIR_temp_val]\n",
    "    train_output = os.path.join(args.out_path, \"train\")\n",
    "    val_output = os.path.join(args.out_path, \"val\")\n",
    "\n",
    "    # Make sure we have a clean parquet space for cudf conversion\n",
    "    for one_path in PREPROCESS_DIR_temp:\n",
    "        if os.path.exists(one_path):\n",
    "            shutil.rmtree(one_path)\n",
    "        os.mkdir(one_path)\n",
    "\n",
    "\n",
    "    ## Get Dask Client\n",
    "\n",
    "    # Deploy a Single-Machine Multi-GPU Cluster\n",
    "    device_size = device_mem_size(kind=\"total\")\n",
    "    cluster = None\n",
    "    if args.protocol == \"ucx\":\n",
    "        UCX_TLS = os.environ.get(\"UCX_TLS\", \"tcp,cuda_copy,cuda_ipc,sockcm\")\n",
    "        os.environ[\"UCX_TLS\"] = UCX_TLS\n",
    "        cluster = LocalCUDACluster(\n",
    "            protocol = args.protocol,\n",
    "            CUDA_VISIBLE_DEVICES = args.devices,\n",
    "            n_workers = len(args.devices.split(\",\")),\n",
    "            enable_nvlink=True,\n",
    "            device_memory_limit = int(device_size * args.device_limit_frac),\n",
    "            dashboard_address=\":\" + args.dashboard_port\n",
    "        )\n",
    "    else:\n",
    "        cluster = LocalCUDACluster(\n",
    "            protocol = args.protocol,\n",
    "            n_workers = len(args.devices.split(\",\")),\n",
    "            CUDA_VISIBLE_DEVICES = args.devices,\n",
    "            device_memory_limit = int(device_size * args.device_limit_frac),\n",
    "            dashboard_address=\":\" + args.dashboard_port\n",
    "        )\n",
    "\n",
    "\n",
    "\n",
    "    # Create the distributed client\n",
    "    client = Client(cluster)\n",
    "    if args.device_pool_frac > 0.01:\n",
    "        setup_rmm_pool(client, int(args.device_pool_frac*device_size))\n",
    "\n",
    "\n",
    "    #calculate the total processing time\n",
    "    runtime = time.time()\n",
    "\n",
    "    #test dataset without the label feature\n",
    "    if args.dataset_type == 'test':\n",
    "        global LABEL_COLUMNS\n",
    "        LABEL_COLUMNS = []\n",
    "\n",
    "    ##-----------------------------------##\n",
    "    # Dask rapids converts txt to parquet\n",
    "    # Dask cudf dataframe = ddf\n",
    "\n",
    "    ## train/valid txt to parquet\n",
    "    train_valid_paths = [(train_input,PREPROCESS_DIR_temp_train),(val_input,PREPROCESS_DIR_temp_val)]\n",
    "\n",
    "    for input, temp_output in train_valid_paths:\n",
    "\n",
    "        ddf = dask_cudf.read_csv(input,sep='\\t',names=LABEL_COLUMNS + CONTINUOUS_COLUMNS + CATEGORICAL_COLUMNS)\n",
    "\n",
    "        ## Convert label col to FP32\n",
    "        if args.parquet_format and args.dataset_type == 'train':\n",
    "            ddf[\"label\"] = ddf['label'].astype('float32')\n",
    "\n",
    "        # Save it as parquet format for better memory usage\n",
    "        ddf.to_parquet(temp_output,header=True)\n",
    "        ##-----------------------------------##\n",
    "\n",
    "    COLUMNS =  LABEL_COLUMNS + CONTINUOUS_COLUMNS + CROSS_COLUMNS + CATEGORICAL_COLUMNS\n",
    "    train_paths = glob.glob(os.path.join(PREPROCESS_DIR_temp_train, \"*.parquet\"))\n",
    "    valid_paths = glob.glob(os.path.join(PREPROCESS_DIR_temp_val, \"*.parquet\"))\n",
    "\n",
    "    categorify_op = Categorify(freq_threshold=args.freq_limit)\n",
    "    cat_features = CATEGORICAL_COLUMNS >> categorify_op\n",
    "    cont_features = CONTINUOUS_COLUMNS >> FillMissing() >> Clip(min_value=0) >> Normalize()\n",
    "    cross_cat_op = Categorify(freq_threshold=args.freq_limit)\n",
    "\n",
    "    features = LABEL_COLUMNS\n",
    "    \n",
    "    if args.criteo_mode == 0:\n",
    "        features += cont_features\n",
    "        if args.feature_cross_list:\n",
    "            feature_pairs = [pair.split(\"_\") for pair in args.feature_cross_list.split(\",\")]\n",
    "            for pair in feature_pairs:\n",
    "                col0 = pair[0]\n",
    "                col1 = pair[1]\n",
    "                features += col0 >> FeatureCross(col1)  >> Rename(postfix=\"_\"+col1) >> cross_cat_op\n",
    "            \n",
    "    features += cat_features\n",
    "\n",
    "    workflow = nvt.Workflow(features, client=client)\n",
    "\n",
    "    logging.info(\"Preprocessing\")\n",
    "\n",
    "    output_format = 'hugectr'\n",
    "    if args.parquet_format:\n",
    "        output_format = 'parquet'\n",
    "\n",
    "    # just for /samples/criteo model\n",
    "    train_ds_iterator = nvt.Dataset(train_paths, engine='parquet', part_size=int(args.part_mem_frac * device_size))\n",
    "    valid_ds_iterator = nvt.Dataset(valid_paths, engine='parquet', part_size=int(args.part_mem_frac * device_size))\n",
    "\n",
    "    shuffle = None\n",
    "    if args.shuffle == \"PER_WORKER\":\n",
    "        shuffle = nvt.io.Shuffle.PER_WORKER\n",
    "    elif args.shuffle == \"PER_PARTITION\":\n",
    "        shuffle = nvt.io.Shuffle.PER_PARTITION\n",
    "\n",
    "    logging.info('Train Datasets Preprocessing.....')\n",
    "\n",
    "    dict_dtypes = {}\n",
    "    for col in CATEGORICAL_COLUMNS:\n",
    "        dict_dtypes[col] = np.int64\n",
    "    if not args.criteo_mode:\n",
    "        for col in CONTINUOUS_COLUMNS:\n",
    "            dict_dtypes[col] = np.float32\n",
    "    for col in CROSS_COLUMNS:\n",
    "        dict_dtypes[col] = np.int64\n",
    "    for col in LABEL_COLUMNS:\n",
    "        dict_dtypes[col] = np.float32\n",
    "    \n",
    "    conts = CONTINUOUS_COLUMNS if not args.criteo_mode else []\n",
    "    \n",
    "    workflow.fit(train_ds_iterator)\n",
    "    \n",
    "    if output_format == 'hugectr':\n",
    "        workflow.transform(train_ds_iterator).to_hugectr(\n",
    "                cats=CATEGORICAL_COLUMNS + CROSS_COLUMNS,\n",
    "                conts=conts,\n",
    "                labels=LABEL_COLUMNS,\n",
    "                output_path=train_output,\n",
    "                shuffle=shuffle,\n",
    "                out_files_per_proc=args.out_files_per_proc,\n",
    "                num_threads=args.num_io_threads)\n",
    "    else:\n",
    "        workflow.transform(train_ds_iterator).to_parquet(\n",
    "                output_path=train_output,\n",
    "                dtypes=dict_dtypes,\n",
    "                cats=CATEGORICAL_COLUMNS + CROSS_COLUMNS,\n",
    "                conts=conts,\n",
    "                labels=LABEL_COLUMNS,\n",
    "                shuffle=shuffle,\n",
    "                out_files_per_proc=args.out_files_per_proc,\n",
    "                num_threads=args.num_io_threads)\n",
    "        \n",
    "        \n",
    "        \n",
    "    ###Getting slot size###    \n",
    "    #--------------------##\n",
    "    embeddings_dict_cat = categorify_op.get_embedding_sizes(CATEGORICAL_COLUMNS)\n",
    "    embeddings_dict_cross = cross_cat_op.get_embedding_sizes(CROSS_COLUMNS)\n",
    "    embeddings = [embeddings_dict_cat[c][0] for c in CATEGORICAL_COLUMNS] + [embeddings_dict_cross[c][0] for c in CROSS_COLUMNS]\n",
    "    \n",
    "    print(embeddings)\n",
    "    ##--------------------##\n",
    "\n",
    "    logging.info('Valid Datasets Preprocessing.....')\n",
    "\n",
    "    if output_format == 'hugectr':\n",
    "        workflow.transform(valid_ds_iterator).to_hugectr(\n",
    "                cats=CATEGORICAL_COLUMNS + CROSS_COLUMNS,\n",
    "                conts=conts,\n",
    "                labels=LABEL_COLUMNS,\n",
    "                output_path=val_output,\n",
    "                shuffle=shuffle,\n",
    "                out_files_per_proc=args.out_files_per_proc,\n",
    "                num_threads=args.num_io_threads)\n",
    "    else:\n",
    "        workflow.transform(valid_ds_iterator).to_parquet(\n",
    "                output_path=val_output,\n",
    "                dtypes=dict_dtypes,\n",
    "                cats=CATEGORICAL_COLUMNS + CROSS_COLUMNS,\n",
    "                conts=conts,\n",
    "                labels=LABEL_COLUMNS,\n",
    "                shuffle=shuffle,\n",
    "                out_files_per_proc=args.out_files_per_proc,\n",
    "                num_threads=args.num_io_threads)\n",
    "\n",
    "    embeddings_dict_cat = categorify_op.get_embedding_sizes(CATEGORICAL_COLUMNS)\n",
    "    embeddings_dict_cross = cross_cat_op.get_embedding_sizes(CROSS_COLUMNS)\n",
    "    embeddings = [embeddings_dict_cat[c][0] for c in CATEGORICAL_COLUMNS] + [embeddings_dict_cross[c][0] for c in CROSS_COLUMNS]\n",
    "    \n",
    "    print(embeddings)\n",
    "    ##--------------------##\n",
    "\n",
    "    ## Shutdown clusters\n",
    "    client.close()\n",
    "    logging.info('NVTabular processing done')\n",
    "\n",
    "    runtime = time.time() - runtime\n",
    "\n",
    "    print(\"\\nDask-NVTabular Criteo Preprocessing\")\n",
    "    print(\"--------------------------------------\")\n",
    "    print(f\"data_path          | {args.data_path}\")\n",
    "    print(f\"output_path        | {args.out_path}\")\n",
    "    print(f\"partition size     | {'%.2f GB'%bytesto(int(args.part_mem_frac * device_size),'g')}\")\n",
    "    print(f\"protocol           | {args.protocol}\")\n",
    "    print(f\"device(s)          | {args.devices}\")\n",
    "    print(f\"rmm-pool-frac      | {(args.device_pool_frac)}\")\n",
    "    print(f\"out-files-per-proc | {args.out_files_per_proc}\")\n",
    "    print(f\"num_io_threads     | {args.num_io_threads}\")\n",
    "    print(f\"shuffle            | {args.shuffle}\")\n",
    "    print(\"======================================\")\n",
    "    print(f\"Runtime[s]         | {runtime}\")\n",
    "    print(\"======================================\\n\")\n",
    "\n",
    "\n",
    "def parse_args():\n",
    "    parser = argparse.ArgumentParser(description=(\"Multi-GPU Criteo Preprocessing\"))\n",
    "\n",
    "    #\n",
    "    # System Options\n",
    "    #\n",
    "\n",
    "    parser.add_argument(\"--data_path\", type=str, help=\"Input dataset path (Required)\")\n",
    "    parser.add_argument(\"--out_path\", type=str, help=\"Directory path to write output (Required)\")\n",
    "    parser.add_argument(\n",
    "        \"-d\",\n",
    "        \"--devices\",\n",
    "        default=os.environ.get(\"CUDA_VISIBLE_DEVICES\", \"0\"),\n",
    "        type=str,\n",
    "        help='Comma-separated list of visible devices (e.g. \"0,1,2,3\"). '\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"-p\",\n",
    "        \"--protocol\",\n",
    "        choices=[\"tcp\", \"ucx\"],\n",
    "        default=\"tcp\",\n",
    "        type=str,\n",
    "        help=\"Communication protocol to use (Default 'tcp')\",\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--device_limit_frac\",\n",
    "        default=0.5,\n",
    "        type=float,\n",
    "        help=\"Worker device-memory limit as a fraction of GPU capacity (Default 0.8). \"\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--device_pool_frac\",\n",
    "        default=0.9,\n",
    "        type=float,\n",
    "        help=\"RMM pool size for each worker  as a fraction of GPU capacity (Default 0.9). \"\n",
    "        \"The RMM pool frac is the same for all GPUs, make sure each one has enough memory size\",\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--num_io_threads\",\n",
    "        default=0,\n",
    "        type=int,\n",
    "        help=\"Number of threads to use when writing output data (Default 0). \"\n",
    "        \"If 0 is specified, multi-threading will not be used for IO.\",\n",
    "    )\n",
    "\n",
    "    #\n",
    "    # Data-Decomposition Parameters\n",
    "    #\n",
    "\n",
    "    parser.add_argument(\n",
    "        \"--part_mem_frac\",\n",
    "        default=0.125,\n",
    "        type=float,\n",
    "        help=\"Maximum size desired for dataset partitions as a fraction \"\n",
    "        \"of GPU capacity (Default 0.125)\",\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--out_files_per_proc\",\n",
    "        default=1,\n",
    "        type=int,\n",
    "        help=\"Number of output files to write on each worker (Default 1)\",\n",
    "    )\n",
    "\n",
    "    #\n",
    "    # Preprocessing Options\n",
    "    #\n",
    "\n",
    "    parser.add_argument(\n",
    "        \"-f\",\n",
    "        \"--freq_limit\",\n",
    "        default=0,\n",
    "        type=int,\n",
    "        help=\"Frequency limit for categorical encoding (Default 0)\",\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"-s\",\n",
    "        \"--shuffle\",\n",
    "        choices=[\"PER_WORKER\", \"PER_PARTITION\", \"NONE\"],\n",
    "        default=\"PER_PARTITION\",\n",
    "        help=\"Shuffle algorithm to use when writing output data to disk (Default PER_PARTITION)\",\n",
    "    )\n",
    "\n",
    "    parser.add_argument(\n",
    "        \"--feature_cross_list\", default=None, type=str, help=\"List of feature crossing cols (e.g. C1_C2, C3_C4)\"\n",
    "    )\n",
    "\n",
    "    #\n",
    "    # Diagnostics Options\n",
    "    #\n",
    "\n",
    "    parser.add_argument(\n",
    "        \"--profile\",\n",
    "        metavar=\"PATH\",\n",
    "        default=None,\n",
    "        type=str,\n",
    "        help=\"Specify a file path to export a Dask profile report (E.g. dask-report.html).\"\n",
    "        \"If this option is excluded from the command, not profile will be exported\",\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--dashboard_port\",\n",
    "        default=\"8787\",\n",
    "        type=str,\n",
    "        help=\"Specify the desired port of Dask's diagnostics-dashboard (Default `3787`). \"\n",
    "        \"The dashboard will be hosted at http://<IP>:<PORT>/status\",\n",
    "    )\n",
    "\n",
    "    #\n",
    "    # Format\n",
    "    #\n",
    "\n",
    "    parser.add_argument('--criteo_mode', type=int, default=0)\n",
    "    parser.add_argument('--parquet_format', type=int, default=1)\n",
    "    parser.add_argument('--dataset_type', type=str, default='train')\n",
    "\n",
    "    args = parser.parse_args()\n",
    "    args.n_workers = len(args.devices.split(\",\"))\n",
    "    return args\n",
    "if __name__ == '__main__':\n",
    "\n",
    "    args = parse_args()\n",
    "\n",
    "    process_NVT(args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f6dbddf4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6f7e259b",
   "metadata": {
    "vscode": {
     "languageId": "shellscript"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2021-07-05 05:41:34,199 NVTabular processing\n",
      "2021-07-05 05:42:00,112 Preprocessing\n",
      "2021-07-05 05:42:00,469 Train Datasets Preprocessing.....\n",
      "[249058, 19561, 14212, 6890, 18592, 4, 6356, 1254, 52, 226170, 80508, 72308, 11, 2169, 7597, 61, 4, 923, 15, 249619, 168974, 243480, 68212, 9169, 75, 34, 278018, 415262]\n",
      "2021-07-05 05:44:17,349 Valid Datasets Preprocessing.....\n",
      "[249058, 19561, 14212, 6890, 18592, 4, 6356, 1254, 52, 226170, 80508, 72308, 11, 2169, 7597, 61, 4, 923, 15, 249619, 168974, 243480, 68212, 9169, 75, 34, 278018, 415262]\n",
      "2021-07-05 05:44:19,138 NVTabular processing done\n",
      "\n",
      "Dask-NVTabular Criteo Preprocessing\n",
      "--------------------------------------\n",
      "data_path          | wdl_train/\n",
      "output_path        | wdl_train/\n",
      "partition size     | 2.77 GB\n",
      "protocol           | tcp\n",
      "device(s)          | 0\n",
      "rmm-pool-frac      | 0.5\n",
      "out-files-per-proc | 1\n",
      "num_io_threads     | 2\n",
      "shuffle            | PER_PARTITION\n",
      "======================================\n",
      "Runtime[s]         | 159.50506210327148\n",
      "======================================\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!python3 /wdl_train/preprocess.py --data_path wdl_train/ \\\n",
    "--out_path wdl_train/ --freq_limit 6 --feature_cross_list C1_C2,C3_C4 \\\n",
    "--device_pool_frac 0.5  --devices '0' --num_io_threads 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4125c1c9",
   "metadata": {},
   "source": [
    "#### Check the preprocessed training data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a8c15e77",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total 14537948\n",
      "-rw-r--r-- 1 root root  3336516608 Jul  5 05:44 0.8870d61b8a1f4deca0f911acfb072999.parquet\n",
      "-rw-r--r-- 1 root root          61 Jul  5 05:44 _file_list.txt\n",
      "-rw-r--r-- 1 root root      602767 Jul  5 05:44 _metadata\n",
      "-rw-r--r-- 1 root root        1538 Jul  5 05:44 _metadata.json\n",
      "drwxr-xr-x 2 root root        4096 Jul  5 05:41 temp-parquet-after-conversion\n",
      "-rwxrwxr-x 1 1025 1025 11549710546 Jul  5 05:39 train.txt\n"
     ]
    }
   ],
   "source": [
    "!ls -ll /wdl_train/train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c321c8a0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>I1</th>\n",
       "      <th>I2</th>\n",
       "      <th>I3</th>\n",
       "      <th>I4</th>\n",
       "      <th>I5</th>\n",
       "      <th>I6</th>\n",
       "      <th>I7</th>\n",
       "      <th>I8</th>\n",
       "      <th>I9</th>\n",
       "      <th>I10</th>\n",
       "      <th>...</th>\n",
       "      <th>C17</th>\n",
       "      <th>C18</th>\n",
       "      <th>C19</th>\n",
       "      <th>C20</th>\n",
       "      <th>C21</th>\n",
       "      <th>C22</th>\n",
       "      <th>C23</th>\n",
       "      <th>C24</th>\n",
       "      <th>C25</th>\n",
       "      <th>C26</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>-0.048792</td>\n",
       "      <td>-0.368150</td>\n",
       "      <td>0.478781</td>\n",
       "      <td>-0.133437</td>\n",
       "      <td>-0.069780</td>\n",
       "      <td>0.068484</td>\n",
       "      <td>0.743047</td>\n",
       "      <td>-0.266159</td>\n",
       "      <td>1.481252</td>\n",
       "      <td>1.386036</td>\n",
       "      <td>...</td>\n",
       "      <td>3</td>\n",
       "      <td>356</td>\n",
       "      <td>10</td>\n",
       "      <td>183947</td>\n",
       "      <td>140830</td>\n",
       "      <td>28449</td>\n",
       "      <td>64057</td>\n",
       "      <td>6432</td>\n",
       "      <td>10</td>\n",
       "      <td>22</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>-0.061206</td>\n",
       "      <td>0.840877</td>\n",
       "      <td>-0.594327</td>\n",
       "      <td>0.148456</td>\n",
       "      <td>-0.209261</td>\n",
       "      <td>-0.206385</td>\n",
       "      <td>-0.064249</td>\n",
       "      <td>-0.281810</td>\n",
       "      <td>-0.760031</td>\n",
       "      <td>-0.470383</td>\n",
       "      <td>...</td>\n",
       "      <td>1</td>\n",
       "      <td>781</td>\n",
       "      <td>10</td>\n",
       "      <td>207893</td>\n",
       "      <td>27876</td>\n",
       "      <td>112273</td>\n",
       "      <td>65971</td>\n",
       "      <td>3414</td>\n",
       "      <td>10</td>\n",
       "      <td>5</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>2 rows × 42 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "         I1        I2        I3        I4        I5        I6        I7  \\\n",
       "0 -0.048792 -0.368150  0.478781 -0.133437 -0.069780  0.068484  0.743047   \n",
       "1 -0.061206  0.840877 -0.594327  0.148456 -0.209261 -0.206385 -0.064249   \n",
       "\n",
       "         I8        I9       I10  ...  C17  C18  C19     C20     C21     C22  \\\n",
       "0 -0.266159  1.481252  1.386036  ...    3  356   10  183947  140830   28449   \n",
       "1 -0.281810 -0.760031 -0.470383  ...    1  781   10  207893   27876  112273   \n",
       "\n",
       "     C23   C24  C25  C26  \n",
       "0  64057  6432   10   22  \n",
       "1  65971  3414   10    5  \n",
       "\n",
       "[2 rows x 42 columns]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "df = pd.read_parquet(\"/wdl_train/train/0.8870d61b8a1f4deca0f911acfb072999.parquet\")\n",
    "df.head(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f671766",
   "metadata": {},
   "source": [
    "### WDL Model Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "aad42841",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting ./model.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile './model.py'\n",
    "import hugectr\n",
    "#from mpi4py import MPI\n",
    "solver = hugectr.CreateSolver(max_eval_batches = 4000,\n",
    "                              batchsize_eval = 2720,\n",
    "                              batchsize = 2720,\n",
    "                              lr = 0.001,\n",
    "                              vvgpu = [[2]],\n",
    "                              repeat_dataset = True,\n",
    "                              i64_input_key = True)\n",
    "\n",
    "reader = hugectr.DataReaderParams(data_reader_type = hugectr.DataReaderType_t.Parquet,\n",
    "                                  source = [\"./train/_file_list.txt\"],\n",
    "                                  eval_source = \"./val/_file_list.txt\",\n",
    "                                  check_type = hugectr.Check_t.Non,\n",
    "                                  slot_size_array = [249058, 19561, 14212, 6890, 18592, 4, 6356, 1254, 52, 226170, 80508, 72308, 11, 2169, 7597, 61, 4, 923, 15, 249619, 168974, 243480, 68212, 9169, 75, 34, 278018, 415262])\n",
    "optimizer = hugectr.CreateOptimizer(optimizer_type = hugectr.Optimizer_t.Adam,\n",
    "                                    update_type = hugectr.Update_t.Global,\n",
    "                                    beta1 = 0.9,\n",
    "                                    beta2 = 0.999,\n",
    "                                    epsilon = 0.0000001)\n",
    "model = hugectr.Model(solver, reader, optimizer)\n",
    "\n",
    "model.add(hugectr.Input(label_dim = 1, label_name = \"label\",\n",
    "                        dense_dim = 13, dense_name = \"dense\",\n",
    "                        data_reader_sparse_param_array = \n",
    "                        [hugectr.DataReaderSparseParam(\"wide_data\", 1, True, 2),\n",
    "                        hugectr.DataReaderSparseParam(\"deep_data\", 2, False, 26)]))\n",
    "\n",
    "model.add(hugectr.SparseEmbedding(embedding_type = hugectr.Embedding_t.DistributedSlotSparseEmbeddingHash, \n",
    "                            workspace_size_per_gpu_in_mb = 24,\n",
    "                            embedding_vec_size = 1,\n",
    "                            combiner = \"sum\",\n",
    "                            sparse_embedding_name = \"sparse_embedding2\",\n",
    "                            bottom_name = \"wide_data\",\n",
    "                            optimizer = optimizer))\n",
    "model.add(hugectr.SparseEmbedding(embedding_type = hugectr.Embedding_t.DistributedSlotSparseEmbeddingHash, \n",
    "                            workspace_size_per_gpu_in_mb = 405,\n",
    "                            embedding_vec_size = 16,\n",
    "                            combiner = \"sum\",\n",
    "                            sparse_embedding_name = \"sparse_embedding1\",\n",
    "                            bottom_name = \"deep_data\",\n",
    "                            optimizer = optimizer))\n",
    "\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.Reshape,\n",
    "                            bottom_names = [\"sparse_embedding1\"],\n",
    "                            top_names = [\"reshape1\"],\n",
    "                            leading_dim=416))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.Reshape,\n",
    "                            bottom_names = [\"sparse_embedding2\"],\n",
    "                            top_names = [\"reshape2\"],\n",
    "                            leading_dim=2))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.ReduceSum,\n",
    "                            bottom_names = [\"reshape2\"],\n",
    "                            top_names = [\"wide_redn\"],\n",
    "                            axis = 1))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.Concat,\n",
    "                            bottom_names = [\"reshape1\", \"dense\"],\n",
    "                            top_names = [\"concat1\"]))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.InnerProduct,\n",
    "                            bottom_names = [\"concat1\"],\n",
    "                            top_names = [\"fc1\"],\n",
    "                            num_output=1024))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.ReLU,\n",
    "                            bottom_names = [\"fc1\"],\n",
    "                            top_names = [\"relu1\"]))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.Dropout,\n",
    "                            bottom_names = [\"relu1\"],\n",
    "                            top_names = [\"dropout1\"],\n",
    "                            dropout_rate=0.5))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.InnerProduct,\n",
    "                            bottom_names = [\"dropout1\"],\n",
    "                            top_names = [\"fc2\"],\n",
    "                            num_output=1024))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.ReLU,\n",
    "                            bottom_names = [\"fc2\"],\n",
    "                            top_names = [\"relu2\"]))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.Dropout,\n",
    "                            bottom_names = [\"relu2\"],\n",
    "                            top_names = [\"dropout2\"],\n",
    "                            dropout_rate=0.5))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.InnerProduct,\n",
    "                            bottom_names = [\"dropout2\"],\n",
    "                            top_names = [\"fc3\"],\n",
    "                            num_output=1))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.Add,\n",
    "                            bottom_names = [\"fc3\", \"wide_redn\"],\n",
    "                            top_names = [\"add1\"]))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.BinaryCrossEntropyLoss,\n",
    "                            bottom_names = [\"add1\", \"label\"],\n",
    "                            top_names = [\"loss\"]))\n",
    "model.compile()\n",
    "model.summary()\n",
    "model.fit(max_iter = 21000, display = 1000, eval_interval = 4000, snapshot = 20000, snapshot_prefix = \"wdl\")\n",
    "model.graph_to_json(graph_config_file = \"wdl.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "39d1f362",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "====================================================Model Init=====================================================\n",
      "[13d04h56m26s][HUGECTR][INFO]: Global seed is 3689394843\n",
      "[13d04h56m31s][HUGECTR][INFO]: Peer-to-peer access cannot be fully enabled.\n",
      "Device 2: A30\n",
      "[13d04h56m31s][HUGECTR][INFO]: num of DataReader workers: 1\n",
      "[13d04h56m31s][HUGECTR][INFO]: max_vocabulary_size_per_gpu_=2097152\n",
      "[13d04h56m31s][HUGECTR][INFO]: max_vocabulary_size_per_gpu_=2211840\n",
      "===================================================Model Compile===================================================\n",
      "[13d04h56m46s][HUGECTR][INFO]: gpu0 start to init embedding\n",
      "[13d04h56m46s][HUGECTR][INFO]: gpu0 init embedding done\n",
      "[13d04h56m46s][HUGECTR][INFO]: gpu0 start to init embedding\n",
      "[13d04h56m46s][HUGECTR][INFO]: gpu0 init embedding done\n",
      "===================================================Model Summary===================================================\n",
      "Label                                   Dense                         Sparse                        \n",
      "label                                   dense                          wide_data,deep_data           \n",
      "(None, 1)                               (None, 13)                              \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "Layer Type                              Input Name                    Output Name                   Output Shape                  \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "DistributedSlotSparseEmbeddingHash      wide_data                     sparse_embedding2             (None, 2, 1)                  \n",
      "DistributedSlotSparseEmbeddingHash      deep_data                     sparse_embedding1             (None, 26, 16)                \n",
      "Reshape                                 sparse_embedding1             reshape1                      (None, 416)                   \n",
      "Reshape                                 sparse_embedding2             reshape2                      (None, 2)                     \n",
      "ReduceSum                               reshape2                      wide_redn                     (None, 1)                     \n",
      "Concat                                  reshape1,dense                concat1                       (None, 429)                   \n",
      "InnerProduct                            concat1                       fc1                           (None, 1024)                  \n",
      "ReLU                                    fc1                           relu1                         (None, 1024)                  \n",
      "Dropout                                 relu1                         dropout1                      (None, 1024)                  \n",
      "InnerProduct                            dropout1                      fc2                           (None, 1024)                  \n",
      "ReLU                                    fc2                           relu2                         (None, 1024)                  \n",
      "Dropout                                 relu2                         dropout2                      (None, 1024)                  \n",
      "InnerProduct                            dropout2                      fc3                           (None, 1)                     \n",
      "Add                                     fc3,wide_redn                 add1                          (None, 1)                     \n",
      "BinaryCrossEntropyLoss                  add1,label                    loss                                                        \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "=====================================================Model Fit=====================================================\n",
      "[13d40h56m46s][HUGECTR][INFO]: Use non-epoch mode with number of iterations: 21000\n",
      "[13d40h56m46s][HUGECTR][INFO]: Training batchsize: 2720, evaluation batchsize: 2720\n",
      "[13d40h56m46s][HUGECTR][INFO]: Evaluation interval: 4000, snapshot interval: 20000\n",
      "[13d40h56m46s][HUGECTR][INFO]: Sparse embedding trainable: 1, dense network trainable: 1\n",
      "[13d40h56m46s][HUGECTR][INFO]: Use mixed precision: 0, scaler: 1.000000, use cuda graph: 1\n",
      "[13d40h56m46s][HUGECTR][INFO]: lr: 0.001000, warmup_steps: 1, decay_start: 0, decay_steps: 1, decay_power: 2.000000, end_lr: 0.000000\n",
      "[13d40h56m46s][HUGECTR][INFO]: Training source file: ./train/_file_list.txt\n",
      "[13d40h56m46s][HUGECTR][INFO]: Evaluation source file: ./val/_file_list.txt\n",
      "[13d40h56m53s][HUGECTR][INFO]: Iter: 1000 Time(1000 iters): 6.920620s Loss: 0.083037 lr:0.001000\n",
      "[13d40h57m00s][HUGECTR][INFO]: Iter: 2000 Time(1000 iters): 6.706585s Loss: 0.120559 lr:0.001000\n",
      "[13d40h57m60s][HUGECTR][INFO]: Iter: 3000 Time(1000 iters): 6.699129s Loss: 0.117169 lr:0.001000\n",
      "[13d40h57m13s][HUGECTR][INFO]: Iter: 4000 Time(1000 iters): 6.758591s Loss: 0.083112 lr:0.001000\n",
      "[13d40h57m23s][HUGECTR][INFO]: Evaluation, AUC: 0.824140\n",
      "[13d40h57m23s][HUGECTR][INFO]: Eval Time for 4000 iters: 9.483117s\n",
      "[13d40h57m29s][HUGECTR][INFO]: Iter: 5000 Time(1000 iters): 16.187022s Loss: 0.131896 lr:0.001000\n",
      "[13d40h57m36s][HUGECTR][INFO]: Iter: 6000 Time(1000 iters): 6.748882s Loss: 0.082966 lr:0.001000\n",
      "[13d40h57m43s][HUGECTR][INFO]: Iter: 7000 Time(1000 iters): 6.761953s Loss: 0.091929 lr:0.001000\n",
      "[13d40h57m50s][HUGECTR][INFO]: Iter: 8000 Time(1000 iters): 6.874048s Loss: 0.080763 lr:0.001000\n",
      "[13d40h57m59s][HUGECTR][INFO]: Evaluation, AUC: 0.826269\n",
      "[13d40h57m59s][HUGECTR][INFO]: Eval Time for 4000 iters: 9.275068s\n",
      "[13d40h58m60s][HUGECTR][INFO]: Iter: 9000 Time(1000 iters): 15.969286s Loss: 0.088093 lr:0.001000\n",
      "[13d40h58m12s][HUGECTR][INFO]: Iter: 10000 Time(1000 iters): 6.652935s Loss: 0.137476 lr:0.001000\n",
      "[13d40h58m19s][HUGECTR][INFO]: Iter: 11000 Time(1000 iters): 6.751184s Loss: 0.116295 lr:0.001000\n",
      "[13d40h58m26s][HUGECTR][INFO]: Iter: 12000 Time(1000 iters): 6.659960s Loss: 0.151319 lr:0.001000\n",
      "[13d40h58m35s][HUGECTR][INFO]: Evaluation, AUC: 0.827362\n",
      "[13d40h58m35s][HUGECTR][INFO]: Eval Time for 4000 iters: 9.378966s\n",
      "[13d40h58m42s][HUGECTR][INFO]: Iter: 13000 Time(1000 iters): 16.001544s Loss: 0.094625 lr:0.001000\n",
      "[13d40h58m48s][HUGECTR][INFO]: Iter: 14000 Time(1000 iters): 6.678430s Loss: 0.121618 lr:0.001000\n",
      "[13d40h58m55s][HUGECTR][INFO]: Iter: 15000 Time(1000 iters): 6.840206s Loss: 0.083302 lr:0.001000\n",
      "[13d40h59m20s][HUGECTR][INFO]: Iter: 16000 Time(1000 iters): 6.489092s Loss: 0.102394 lr:0.001000\n",
      "[13d40h59m11s][HUGECTR][INFO]: Evaluation, AUC: 0.829899\n",
      "[13d40h59m11s][HUGECTR][INFO]: Eval Time for 4000 iters: 9.338721s\n",
      "[13d40h59m18s][HUGECTR][INFO]: Iter: 17000 Time(1000 iters): 15.868251s Loss: 0.108997 lr:0.001000\n",
      "[13d40h59m24s][HUGECTR][INFO]: Iter: 18000 Time(1000 iters): 5.960831s Loss: 0.098293 lr:0.001000\n",
      "[13d40h59m29s][HUGECTR][INFO]: Iter: 19000 Time(1000 iters): 5.980448s Loss: 0.071080 lr:0.001000\n",
      "[13d40h59m35s][HUGECTR][INFO]: Iter: 20000 Time(1000 iters): 5.984280s Loss: 0.115342 lr:0.001000\n",
      "[13d40h59m45s][HUGECTR][INFO]: Evaluation, AUC: 0.828875\n",
      "[13d40h59m45s][HUGECTR][INFO]: Eval Time for 4000 iters: 9.337684s\n",
      "[13d40h59m45s][HUGECTR][INFO]: Rank0: Write hash table to file\n",
      "[13d40h59m45s][HUGECTR][INFO]: Rank0: Write hash table to file\n",
      "[13d40h59m45s][HUGECTR][INFO]: Dumping sparse weights to files, successful\n",
      "[13d40h59m45s][HUGECTR][INFO]: Rank0: Write optimzer state to file\n",
      "[13d40h59m45s][HUGECTR][INFO]: Done\n",
      "[13d40h59m45s][HUGECTR][INFO]: Rank0: Write optimzer state to file\n",
      "[13d40h59m45s][HUGECTR][INFO]: Done\n",
      "[13d40h59m45s][HUGECTR][INFO]: Rank0: Write optimzer state to file\n",
      "[13d40h59m45s][HUGECTR][INFO]: Done\n",
      "[13d40h59m45s][HUGECTR][INFO]: Rank0: Write optimzer state to file\n",
      "[13d40h59m45s][HUGECTR][INFO]: Done\n",
      "[13d40h59m45s][HUGECTR][INFO]: Dumping sparse optimzer states to files, successful\n",
      "[13d40h59m45s][HUGECTR][INFO]: Dumping dense weights to file, successful\n",
      "[13d40h59m45s][HUGECTR][INFO]: Dumping dense optimizer states to file, successful\n",
      "[13d40h59m45s][HUGECTR][INFO]: Dumping untrainable weights to file, successful\n",
      "[13d40h59m51s][HUGECTR][INFO]: Save the model graph to wdl.json, successful\n"
     ]
    }
   ],
   "source": [
    "!python ./model.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "6004ea29",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total 301620\n",
      "-rw-rw-r-- 1 1025 1025     49824 Jul  6 07:00 HugeCTR_WDL_Training.ipynb\n",
      "drwxr-xr-x 2 root root      4096 Jul  5 05:43 categories\n",
      "drwxr-xr-x 3 root root      4096 Jul  5 05:44 dask-worker-space\n",
      "-rw-r--r-- 1 root root      5539 Jul  6 07:00 model.py\n",
      "-rw-r--r-- 1 root root     14265 Jul  6 06:59 preprocess.py\n",
      "drwxr-xr-x 3 root root      4096 Jul  5 23:34 train\n",
      "drwxr-xr-x 3 root root      4096 Jul  5 05:44 val\n",
      "-rw-r--r-- 1 root root  17108704 Jul  6 03:28 wdl0_opt_sparse_20000.model\n",
      "drwxr-xr-x 2 root root      4096 Jul  5 06:32 wdl0_sparse_20000.model\n",
      "-rw-r--r-- 1 root root 273739264 Jul  6 03:28 wdl1_opt_sparse_20000.model\n",
      "drwxr-xr-x 2 root root      4096 Jul  5 06:32 wdl1_sparse_20000.model\n",
      "-rw-r--r-- 1 root root   5963780 Jul  6 03:28 wdl_dense_20000.model\n",
      "-rw-r--r-- 1 root root      3158 Jul  6 03:28 wdl_infer.json\n",
      "-rw-r--r-- 1 root root  11927560 Jul  6 03:28 wdl_opt_dense_20000.model\n"
     ]
    }
   ],
   "source": [
    "!ls -ll"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "da964708",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['/wdl_infer/model/wdl/1/wdl0_sparse_20000.model/', '/wdl_infer/model/wdl/1/wdl1_sparse_20000.model']\n",
      "[14d04h23m54s][HUGECTR][INFO]: default_emb_vec_value is not specified using default: 0.000000\n",
      "[14d04h23m54s][HUGECTR][INFO]: default_emb_vec_value is not specified using default: 0.000000\n",
      "Local RocksDB is initializing the embedding table: wdl0\n",
      "Last Iteration insert successfully\n",
      "Local RocksDB is initializing the embedding table: wdl1\n",
      "Last Iteration insert successfully\n",
      "[14d04h24m08s][HUGECTR][INFO]: Global seed is 2483322206\n",
      "[14d04h24m08s][HUGECTR][INFO]: Device to NUMA mapping:\n",
      "  GPU 2 ->  node 1\n",
      "\n",
      "[14d04h24m13s][HUGECTR][INFO]: Peer-to-peer access cannot be fully enabled.\n",
      "[14d04h24m13s][HUGECTR][INFO]: Start all2all warmup\n",
      "[14d04h24m13s][HUGECTR][INFO]: End all2all warmup\n",
      "[14d04h24m13s][HUGECTR][INFO]: Use mixed precision: 0\n",
      "[14d04h24m13s][HUGECTR][INFO]: start create embedding for inference\n",
      "[14d04h24m13s][HUGECTR][INFO]: sparse_input name wide_data\n",
      "[14d04h24m13s][HUGECTR][INFO]: sparse_input name deep_data\n",
      "[14d04h24m13s][HUGECTR][INFO]: create embedding for inference success\n",
      "[14d04h24m13s][HUGECTR][INFO]: Inference stage skip BinaryCrossEntropyLoss layer, replaced by Sigmoid layer\n",
      "Rocksdb gets missing keys from model: wdl and table: 0\n",
      "Rocksdb gets missing keys from model: wdl and table: 1\n",
      "WDL multi-embedding table inference result is [0.2726621925830841, 0.16786302626132965, 0.06844793260097504, 0.21687281131744385, 0.28839486837387085, 0.09961184859275818, 0.1451544463634491, 0.1859627217054367, 0.1754387617111206, 0.14994166791439056]\n",
      "[HUGECTR][INFO] WDL multi-embedding table inference using GPU cache, prediction error is less  than threshold:0.0001, error is 1.1102230246251565e-16\n"
     ]
    }
   ],
   "source": [
    "!python /wdl_infer/wdl_python_infer.py \"wdl\" \"/wdl_infer/model/wdl/1/wdl.json\" \\\n",
    "\"/wdl_infer/model/wdl/1/wdl_dense_20000.model\" \\\n",
    "\"/wdl_infer/model/wdl/1/wdl0_sparse_20000.model/,/wdl_infer/model/wdl/1/wdl1_sparse_20000.model\" \\\n",
    "\"/wdl_infer/first_ten.csv\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b649a6a9",
   "metadata": {},
   "source": [
    "## Prepare Inference Request"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "130dd6ea",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total 637376\n",
      "-rw-r--r-- 1 root root 142856977 Jul  5 05:44 0.110d099942694a5cbf1b71eb73e10f27.parquet\n",
      "-rw-r--r-- 1 root root        51 Jul  6 07:02 _file_list.txt\n",
      "-rw-r--r-- 1 root root     27701 Jul  5 05:44 _metadata\n",
      "-rw-r--r-- 1 root root      1537 Jul  5 05:44 _metadata.json\n",
      "drwxr-xr-x 2 root root      4096 Jul  5 05:42 temp-parquet-after-conversion\n",
      "-rw-r--r-- 1 1025 1025 509766965 Jul  5 04:45 test.txt\n"
     ]
    }
   ],
   "source": [
    "!ls -l /wdl_train/val"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "449cb610",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>I1</th>\n",
       "      <th>I2</th>\n",
       "      <th>I3</th>\n",
       "      <th>I4</th>\n",
       "      <th>I5</th>\n",
       "      <th>I6</th>\n",
       "      <th>I7</th>\n",
       "      <th>I8</th>\n",
       "      <th>I9</th>\n",
       "      <th>I10</th>\n",
       "      <th>...</th>\n",
       "      <th>C17</th>\n",
       "      <th>C18</th>\n",
       "      <th>C19</th>\n",
       "      <th>C20</th>\n",
       "      <th>C21</th>\n",
       "      <th>C22</th>\n",
       "      <th>C23</th>\n",
       "      <th>C24</th>\n",
       "      <th>C25</th>\n",
       "      <th>C26</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0.061161</td>\n",
       "      <td>0.974006</td>\n",
       "      <td>-0.594327</td>\n",
       "      <td>-0.157301</td>\n",
       "      <td>-0.224758</td>\n",
       "      <td>0.618222</td>\n",
       "      <td>-0.064249</td>\n",
       "      <td>-0.281810</td>\n",
       "      <td>-0.760031</td>\n",
       "      <td>1.386036</td>\n",
       "      <td>...</td>\n",
       "      <td>2</td>\n",
       "      <td>666</td>\n",
       "      <td>1</td>\n",
       "      <td>33722</td>\n",
       "      <td>24373</td>\n",
       "      <td>91481</td>\n",
       "      <td>62242</td>\n",
       "      <td>7673</td>\n",
       "      <td>44</td>\n",
       "      <td>28</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>-0.061206</td>\n",
       "      <td>-0.437431</td>\n",
       "      <td>0.156849</td>\n",
       "      <td>-0.146861</td>\n",
       "      <td>-0.193763</td>\n",
       "      <td>0.893091</td>\n",
       "      <td>-0.064249</td>\n",
       "      <td>0.286841</td>\n",
       "      <td>-0.109336</td>\n",
       "      <td>3.242455</td>\n",
       "      <td>...</td>\n",
       "      <td>1</td>\n",
       "      <td>666</td>\n",
       "      <td>10</td>\n",
       "      <td>0</td>\n",
       "      <td>97438</td>\n",
       "      <td>0</td>\n",
       "      <td>21446</td>\n",
       "      <td>4472</td>\n",
       "      <td>56</td>\n",
       "      <td>19</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0.043427</td>\n",
       "      <td>-0.464600</td>\n",
       "      <td>-0.379705</td>\n",
       "      <td>-0.120014</td>\n",
       "      <td>0.054203</td>\n",
       "      <td>-0.206385</td>\n",
       "      <td>-0.064249</td>\n",
       "      <td>-0.093999</td>\n",
       "      <td>-0.543133</td>\n",
       "      <td>-0.470383</td>\n",
       "      <td>...</td>\n",
       "      <td>1</td>\n",
       "      <td>575</td>\n",
       "      <td>10</td>\n",
       "      <td>0</td>\n",
       "      <td>46601</td>\n",
       "      <td>0</td>\n",
       "      <td>12090</td>\n",
       "      <td>540</td>\n",
       "      <td>10</td>\n",
       "      <td>17</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>-0.059432</td>\n",
       "      <td>-0.273058</td>\n",
       "      <td>-0.487016</td>\n",
       "      <td>-0.143878</td>\n",
       "      <td>-0.193763</td>\n",
       "      <td>-0.206385</td>\n",
       "      <td>-0.064249</td>\n",
       "      <td>-0.279201</td>\n",
       "      <td>-0.109336</td>\n",
       "      <td>-0.470383</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>351</td>\n",
       "      <td>10</td>\n",
       "      <td>125237</td>\n",
       "      <td>4329</td>\n",
       "      <td>238309</td>\n",
       "      <td>0</td>\n",
       "      <td>8488</td>\n",
       "      <td>56</td>\n",
       "      <td>22</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>-0.048792</td>\n",
       "      <td>-0.418412</td>\n",
       "      <td>0.693403</td>\n",
       "      <td>0.300589</td>\n",
       "      <td>-0.193763</td>\n",
       "      <td>-0.206385</td>\n",
       "      <td>-0.064249</td>\n",
       "      <td>-0.281810</td>\n",
       "      <td>0.902856</td>\n",
       "      <td>-0.470383</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>575</td>\n",
       "      <td>7</td>\n",
       "      <td>69747</td>\n",
       "      <td>76381</td>\n",
       "      <td>207280</td>\n",
       "      <td>0</td>\n",
       "      <td>444</td>\n",
       "      <td>73</td>\n",
       "      <td>22</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 42 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "         I1        I2        I3        I4        I5        I6        I7  \\\n",
       "0  0.061161  0.974006 -0.594327 -0.157301 -0.224758  0.618222 -0.064249   \n",
       "1 -0.061206 -0.437431  0.156849 -0.146861 -0.193763  0.893091 -0.064249   \n",
       "2  0.043427 -0.464600 -0.379705 -0.120014  0.054203 -0.206385 -0.064249   \n",
       "3 -0.059432 -0.273058 -0.487016 -0.143878 -0.193763 -0.206385 -0.064249   \n",
       "4 -0.048792 -0.418412  0.693403  0.300589 -0.193763 -0.206385 -0.064249   \n",
       "\n",
       "         I8        I9       I10  ...  C17  C18  C19     C20    C21     C22  \\\n",
       "0 -0.281810 -0.760031  1.386036  ...    2  666    1   33722  24373   91481   \n",
       "1  0.286841 -0.109336  3.242455  ...    1  666   10       0  97438       0   \n",
       "2 -0.093999 -0.543133 -0.470383  ...    1  575   10       0  46601       0   \n",
       "3 -0.279201 -0.109336 -0.470383  ...    0  351   10  125237   4329  238309   \n",
       "4 -0.281810  0.902856 -0.470383  ...    0  575    7   69747  76381  207280   \n",
       "\n",
       "     C23   C24  C25  C26  \n",
       "0  62242  7673   44   28  \n",
       "1  21446  4472   56   19  \n",
       "2  12090   540   10   17  \n",
       "3      0  8488   56   22  \n",
       "4      0   444   73   22  \n",
       "\n",
       "[5 rows x 42 columns]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "df = pd.read_parquet(\"/wdl_train/val/0.110d099942694a5cbf1b71eb73e10f27.parquet\")\n",
    "\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "85069108",
   "metadata": {},
   "outputs": [],
   "source": [
    "df.head(10).to_csv('/wdl_train/infer_test.csv', sep=',', index=False,header=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87f219ba",
   "metadata": {},
   "source": [
    "## Create prediction scripts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e98d68e7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting /wdl_train/wdl_predict.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile '/wdl_train/wdl_predict.py'\n",
    "from hugectr.inference import InferenceParams, CreateInferenceSession\n",
    "import hugectr\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import sys\n",
    "from mpi4py import MPI\n",
    "def wdl_inference(model_name, network_file, dense_file, embedding_file_list, data_file,enable_cache,dbtype=hugectr.Database_t.Local,rocksdb_path=\"\"):\n",
    "    CATEGORICAL_COLUMNS=[\"C\" + str(x) for x in range(1, 27)]+[\"C1_C2\",\"C3_C4\"]\n",
    "    CONTINUOUS_COLUMNS=[\"I\" + str(x) for x in range(1, 14)]\n",
    "    LABEL_COLUMNS = ['label']\n",
    "    emb_size = [249058, 19561, 14212, 6890, 18592, 4, 6356, 1254, 52, 226170, 80508, 72308, 11, 2169, 7597, 61, 4, 923, 15, 249619, 168974, 243480, 68212, 9169, 75, 34, 278018, 415262]\n",
    "    shift = np.insert(np.cumsum(emb_size), 0, 0)[:-1]\n",
    "    test_df=pd.read_csv(data_file,sep=',')\n",
    "    config_file = network_file\n",
    "    row_ptrs = list(range(0,21))+list(range(0,261))\n",
    "    dense_features =  list(test_df[CONTINUOUS_COLUMNS].values.flatten())\n",
    "    test_df[CATEGORICAL_COLUMNS].astype(np.int64)\n",
    "    embedding_columns = list((test_df[CATEGORICAL_COLUMNS]+shift).values.flatten())\n",
    "    \n",
    "\n",
    "    # create parameter server, embedding cache and inference session\n",
    "    inference_params = InferenceParams(model_name = model_name,\n",
    "                                max_batchsize = 64,\n",
    "                                hit_rate_threshold = 0.5,\n",
    "                                dense_model_file = dense_file,\n",
    "                                sparse_model_files = embedding_file_list,\n",
    "                                device_id = 2,\n",
    "                                use_gpu_embedding_cache = enable_cache,\n",
    "                                cache_size_percentage = 0.9,\n",
    "                                i64_input_key = True,\n",
    "                                use_mixed_precision = False,\n",
    "                                db_type = dbtype,\n",
    "                                rocksdb_path=rocksdb_path,\n",
    "                                cache_size_percentage_redis=0.5)\n",
    "    inference_session = CreateInferenceSession(config_file, inference_params)\n",
    "    output = inference_session.predict(dense_features, embedding_columns, row_ptrs)\n",
    "    print(\"WDL multi-embedding table inference result is {}\".format(output))\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    model_name = sys.argv[1]\n",
    "    print(\"{} multi-embedding table prediction\".format(model_name))\n",
    "    network_file = sys.argv[2]\n",
    "    print(\"{} multi-embedding table prediction network is {}\".format(model_name,network_file))\n",
    "    dense_file = sys.argv[3]\n",
    "    print(\"{} multi-embedding table prediction dense file is {}\".format(model_name,dense_file))\n",
    "    embedding_file_list = str(sys.argv[4]).split(',')\n",
    "    print(\"{} multi-embedding table prediction sparse files are {}\".format(model_name,embedding_file_list))\n",
    "    data_file = sys.argv[5]\n",
    "    print(\"{} multi-embedding table prediction input data path is {}\".format(model_name,data_file))\n",
    "    input_dbtype = sys.argv[6]\n",
    "    print(\"{} multi-embedding table prediction input dbtype path is {}\".format(model_name,input_dbtype))\n",
    "    if input_dbtype==\"local\":\n",
    "        wdl_inference(model_name, network_file, dense_file, embedding_file_list, data_file, True, hugectr.Database_t.Local)\n",
    "    if input_dbtype==\"rocksdb\":\n",
    "        rocksdb_path = sys.argv[7]\n",
    "        print(\"{} multi-embedding table prediction rocksdb_path path is {}\".format(model_name,rocksdb_path))\n",
    "        wdl_inference(model_name, network_file, dense_file, embedding_file_list, data_file, True, hugectr.Database_t.RocksDB,rocksdb_path)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f58bdd0f",
   "metadata": {},
   "source": [
    "## Prediction\n",
    "\n",
    "Use different types of databases as a local parameter server to get the wide and deep model prediction results."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "478b1f54",
   "metadata": {},
   "source": [
    "### Load model embedding tables into local memory as parameter server"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "604cb6f6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "wdl multi-embedding table prediction\n",
      "wdl multi-embedding table prediction network is /wdl_infer/model/wdl/1/wdl.json\n",
      "wdl multi-embedding table prediction dense file is /wdl_infer/model/wdl/1/wdl_dense_20000.model\n",
      "wdl multi-embedding table prediction sparse files are ['/wdl_infer/model/wdl/1/wdl0_sparse_20000.model/', '/wdl_infer/model/wdl/1/wdl1_sparse_20000.model']\n",
      "wdl multi-embedding table prediction input data path is /wdl_train/infer_test.csv\n",
      "wdl multi-embedding table prediction input dbtype path is local\n",
      "[15d09h17m26s][HUGECTR][INFO]: default_emb_vec_value is not specified using default: 0.000000\n",
      "[15d09h17m26s][HUGECTR][INFO]: default_emb_vec_value is not specified using default: 0.000000\n",
      "[15d09h17m32s][HUGECTR][INFO]: Global seed is 1361897547\n",
      "[15d09h17m32s][HUGECTR][INFO]: Device to NUMA mapping:\n",
      "  GPU 2 ->  node 1\n",
      "\n",
      "[15d09h17m48s][HUGECTR][INFO]: Peer-to-peer access cannot be fully enabled.\n",
      "[15d09h17m48s][HUGECTR][INFO]: Start all2all warmup\n",
      "[15d09h17m48s][HUGECTR][INFO]: End all2all warmup\n",
      "[15d09h17m48s][HUGECTR][INFO]: Use mixed precision: 0\n",
      "[15d09h17m48s][HUGECTR][INFO]: start create embedding for inference\n",
      "[15d09h17m48s][HUGECTR][INFO]: sparse_input name wide_data\n",
      "[15d09h17m48s][HUGECTR][INFO]: sparse_input name deep_data\n",
      "[15d09h17m48s][HUGECTR][INFO]: create embedding for inference success\n",
      "[15d09h17m48s][HUGECTR][INFO]: Inference stage skip BinaryCrossEntropyLoss layer, replaced by Sigmoid layer\n",
      "WDL multi-embedding table inference result is [0.019959857687354088, 0.025274723768234253, 0.017903145402669907, 0.006932722870260477, 0.02339070290327072, 0.022747302427887917, 0.05989734083414078, 0.015981541946530342, 0.005822415463626385, 0.01423134095966816]\n"
     ]
    }
   ],
   "source": [
    "!python /wdl_train/wdl_predict.py \"wdl\" \"/wdl_infer/model/wdl/1/wdl.json\" \"/wdl_infer/model/wdl/1/wdl_dense_20000.model\" \"/wdl_infer/model/wdl/1/wdl0_sparse_20000.model/,/wdl_infer/model/wdl/1/wdl1_sparse_20000.model\" \"/wdl_train/infer_test.csv\" \"local\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7284260e",
   "metadata": {},
   "source": [
    "### Load model embedding tables into local RocksDB as a parameter Server\n",
    "\n",
    "Create a RocksDB directory with read and write permissions for storing model embedded tables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b9f750d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p -m 700 /wdl_train/rocksdb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c59e4b05",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "wdl multi-embedding table prediction\n",
      "wdl multi-embedding table prediction network is /wdl_infer/model/wdl/1/wdl.json\n",
      "wdl multi-embedding table prediction dense file is /wdl_infer/model/wdl/1/wdl_dense_20000.model\n",
      "wdl multi-embedding table prediction sparse files are ['/wdl_infer/model/wdl/1/wdl0_sparse_20000.model/', '/wdl_infer/model/wdl/1/wdl1_sparse_20000.model']\n",
      "wdl multi-embedding table prediction input data path is /wdl_train/infer_test.csv\n",
      "wdl multi-embedding table prediction input dbtype path is rocksdb\n",
      "wdl multi-embedding table prediction rocksdb_path path is /wdl_train/rocksdb\n",
      "[15d12h32m00s][HUGECTR][INFO]: default_emb_vec_value is not specified using default: 0.000000\n",
      "[15d12h32m00s][HUGECTR][INFO]: default_emb_vec_value is not specified using default: 0.000000\n",
      "Local RocksDB is initializing the embedding table: wdl0\n",
      "Last Iteration insert successfully\n",
      "Local RocksDB is initializing the embedding table: wdl1\n",
      "Last Iteration insert successfully\n",
      "[15d12h32m07s][HUGECTR][INFO]: Global seed is 1156574989\n",
      "[15d12h32m08s][HUGECTR][INFO]: Device to NUMA mapping:\n",
      "  GPU 2 ->  node 1\n",
      "\n",
      "[15d12h32m21s][HUGECTR][INFO]: Peer-to-peer access cannot be fully enabled.\n",
      "[15d12h32m21s][HUGECTR][INFO]: Start all2all warmup\n",
      "[15d12h32m21s][HUGECTR][INFO]: End all2all warmup\n",
      "[15d12h32m21s][HUGECTR][INFO]: Use mixed precision: 0\n",
      "[15d12h32m21s][HUGECTR][INFO]: start create embedding for inference\n",
      "[15d12h32m21s][HUGECTR][INFO]: sparse_input name wide_data\n",
      "[15d12h32m21s][HUGECTR][INFO]: sparse_input name deep_data\n",
      "[15d12h32m21s][HUGECTR][INFO]: create embedding for inference success\n",
      "[15d12h32m21s][HUGECTR][INFO]: Inference stage skip BinaryCrossEntropyLoss layer, replaced by Sigmoid layer\n",
      "Rocksdb gets missing keys from model: wdl and table: 0\n",
      "Rocksdb gets missing keys from model: wdl and table: 1\n",
      "WDL multi-embedding table inference result is [0.019959857687354088, 0.025274723768234253, 0.017903145402669907, 0.006932722870260477, 0.02339070290327072, 0.022747302427887917, 0.05989734083414078, 0.015981541946530342, 0.005822415463626385, 0.01423134095966816]\n"
     ]
    }
   ],
   "source": [
    "!python /wdl_train/wdl_predict.py \"wdl\" \"/wdl_infer/model/wdl/1/wdl.json\" \\\n",
    "\"/wdl_infer/model/wdl/1/wdl_dense_20000.model\" \\\n",
    "\"/wdl_infer/model/wdl/1/wdl0_sparse_20000.model/,/wdl_infer/model/wdl/1/wdl1_sparse_20000.model\" \\\n",
    "\"/wdl_train/infer_test.csv\" \\\n",
    "\"rocksdb\"  \"/wdl_train/rocksdb\""
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
